{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import time\n",
    "\n",
    "# Define the neural network model\n",
    "class FullyConnectedNN(torch.nn.Module):\n",
    "    def __init__(self, input_dim, hidden_dim, num_layers):\n",
    "        super(FullyConnectedNN, self).__init__()\n",
    "        self.input_dim = input_dim\n",
    "        self.hidden_dim = hidden_dim\n",
    "        self.num_layers = num_layers\n",
    "        self.layers = torch.nn.ModuleList()\n",
    "        self.layers.append(torch.nn.Linear(input_dim, hidden_dim))\n",
    "        for _ in range(num_layers - 1):\n",
    "            self.layers.append(torch.nn.Linear(hidden_dim, hidden_dim))\n",
    "    \n",
    "    def forward(self, x):\n",
    "        for layer in self.layers:\n",
    "            x = torch.tanh(layer(x))\n",
    "        return x\n",
    "\n",
    "# # Define the function u_theta\n",
    "# def u_theta(x, model):\n",
    "#     return model(x)\n",
    "\n",
    "# Define a function to compute the time for u_theta, nabla u_theta, and Delta u_theta\n",
    "def compute_times(model, input_dim, num_samples=1000):\n",
    "    device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
    "    model.to(device)\n",
    "    \n",
    "    u_theta_time       = 0\n",
    "    nabla_u_theta_time = 0\n",
    "    delta_u_theta_time = 0\n",
    "    for step in range(2000):\n",
    "        # Randomly generate sample points\n",
    "        x = torch.randn(num_samples, input_dim).to(device).requires_grad_(True)\n",
    "        \n",
    "        # Compute u_theta and measure the time\n",
    "        start_time = time.time()\n",
    "        u = model(x)\n",
    "        if step >= 1000:\n",
    "            u_theta_time += time.time() - start_time\n",
    "\n",
    "        # Compute nabla u_theta and measure the time\n",
    "        start_time = time.time()\n",
    "        du = torch.autograd.grad(u, x, \n",
    "                                grad_outputs=torch.ones_like(u), \n",
    "                                create_graph=True, \n",
    "                                retain_graph=True)[0]\n",
    "        if step >= 1000:\n",
    "            nabla_u_theta_time += time.time() - start_time\n",
    "\n",
    "        # Compute Delta u_theta and measure the time\n",
    "        start_time = time.time()\n",
    "        laplace = torch.zeros_like(u)\n",
    "        for i in range(input_dim):\n",
    "            d2u = torch.autograd.grad(du[:, i], x, \n",
    "                                    grad_outputs=torch.ones_like(du[:, i]), \n",
    "                                    create_graph=True, \n",
    "                                    retain_graph=True)[0][:, i]\n",
    "            laplace += d2u.reshape(-1, 1)\n",
    "        if step >= 1000:\n",
    "            delta_u_theta_time += time.time() - start_time\n",
    "\n",
    "    return u_theta_time, nabla_u_theta_time, delta_u_theta_time\n",
    "\n",
    "# Loop through different input dimensions and compute the times\n",
    "for input_dim in [1, 2, 10, 50, 100, 200]:\n",
    "    print(\"============= input_dim\", input_dim, \"==============\")\n",
    "    # Parameter settings\n",
    "    hidden_dim = 200\n",
    "    num_layers = 6\n",
    "\n",
    "    # Create the neural network model\n",
    "    model = FullyConnectedNN(input_dim, hidden_dim, num_layers)\n",
    "\n",
    "    # Compute the times\n",
    "    u_theta_time, nabla_u_theta_time, delta_u_theta_time = compute_times(model, input_dim)\n",
    "\n",
    "    print(\"Computing time for u_theta:\", u_theta_time)\n",
    "    print(\"Computing time for nabla u_theta:\", nabla_u_theta_time)\n",
    "    print(\"Computing time for Delta u_theta:\", delta_u_theta_time)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "torch38",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
